package middleware

import (
	"bp/conf"
	"bp/internal/util"
	"context"
	"crypto/rsa"
	"fmt"
	"github.com/dgrijalva/jwt-go"
	"go.uber.org/zap"
	"net/http"
	"sync"
)

const (
	authorizationTokenHeadKey = "Authorization"
)

var (
	authorizationPubKey  *rsa.PublicKey
	authorizationPubOnce sync.Once
)

type (
	authClaims struct {
		jwt.StandardClaims
		*util.AuthorizationToken
	}
)

// WithAuthorizationToken 提取请求的token
func WithAuthorizationToken(slog *zap.SugaredLogger, vars *conf.VarsSection) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			ctx := r.Context()
			tokenStr := util.ExtractToken(r.Header.Get(authorizationTokenHeadKey))
			if tokenStr == "" {
				http.Error(w, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
				return
			}
			tokenPtr, err := VerifyAuthorizationToken(tokenStr, vars.AuthorizationPubKey)
			if err != nil {
				slog.Errorf("验证token err:%#v", err.Error())
				http.Error(w, "验证token失败", http.StatusUnauthorized)
				return
			}
			slog.Infof("token got")
			ctx = context.WithValue(ctx, util.AuthorizationTokenKey{}, tokenPtr)
			r = r.WithContext(ctx)
			
			next.ServeHTTP(w, r)
		})
	}
}

// LoadAuthorizationPubKey 加载公钥
func LoadAuthorizationPubKey(key string) func() {
	return func() {
		authorizationPubKey, _ = jwt.ParseRSAPublicKeyFromPEM([]byte(key))
	}
}

// VerifyAuthorizationToken 比邻C端 token 校验
func VerifyAuthorizationToken(tokenStr, key string) (atoken *util.AuthorizationToken, err error) {
	authorizationPubOnce.Do(LoadAuthorizationPubKey(key))

	parse := &jwt.Parser{
		ValidMethods:         nil,
		UseJSONNumber:        false,
		SkipClaimsValidation: true,
	}
	token, err := parse.ParseWithClaims(tokenStr, &authClaims{}, func(token *jwt.Token) (i interface{}, e error) {
		return authorizationPubKey, nil
	})
	if err != nil {
		return
	}
	claims := token.Claims.(*authClaims)
	err = claims.Valid()
	if err != nil {
		return
	}
	atoken = claims.AuthorizationToken
	if atoken == nil {
		err = fmt.Errorf("jwt payload 为空")
		return
	}
	return
}
